package com.shujia.streeam

import kafka.common.TopicAndPartition
import org.apache.hadoop.hbase.{HBaseConfiguration, HColumnDescriptor, HTableDescriptor}
import org.apache.hadoop.hbase.client.{HBaseAdmin, HConnectionManager, Put}
import org.apache.hadoop.hbase.client.Scan
import org.apache.hadoop.hbase.util.Bytes

object OffsetToHbaseUtil {

  def main(args: Array[String]): Unit = {
    createTable("111")
    println(getOffset("111", "student", 3))

  }


  val configuration = HBaseConfiguration.create()
  configuration.set("hbase.zookeeper.quorum", "node2:2181,node3:2181.node4:2181");
  val connection = HConnectionManager.createConnection(configuration)
  //对表操作时使用
  val hBaseAdmin = new HBaseAdmin(configuration)

  def createTable(groupId: String): Unit = {
    val tableName = "offset_" + groupId

    if (hBaseAdmin.tableExists(tableName)) {
      return
    }

    val desc = new HTableDescriptor(tableName)
    desc.addFamily(new HColumnDescriptor("info"))
    hBaseAdmin.createTable(desc)
  }


  /**
    * 获取初始偏移量
    * 首先去hbase中获取，如果hbase中没有  返回默认值 从头开始消费
    */
  def getOffset(groupId: String, topic: String, partitions: Int): Map[TopicAndPartition, Long] = {

    val map = scala.collection.mutable.Map[TopicAndPartition, Long]()

    val tableName = "offset_" + groupId

    val table = connection.getTable(tableName)

    val scan = new Scan()

    val resultScanner = table.getScanner(scan)

    var next = resultScanner.next()

    while (next != null) {
      val rowkey = Bytes.toString(next.getRow)
      val split = rowkey.split("-")
      val topic = split(0)
      val partition = split(1)
      val tAndP = TopicAndPartition(topic, partition.toInt)
      val fromOffset = Bytes.toString(next.getValue("info".getBytes, "offset".getBytes())).toLong
      map.put(tAndP, fromOffset)
      next = resultScanner.next()
    }

    //如果hbase中没有保存着groupId 的偏移量  那么久从头消费
    if (map.isEmpty) {
      for (partition <- 0.until(partitions)) {
        map.put(TopicAndPartition(topic, partition), 0L)
      }
    }

    map.toMap

  }

  //更新偏移量
  def updataOffset(groupId: String, topic: String, partition: Int, untilOffset: Long): Unit = {

    val tableName = "offset_" + groupId

    val table = connection.getTable(tableName)

    val rowkey = topic + "-" + partition

    val put = new Put(rowkey.getBytes())

    put.add("info".getBytes(), "offset".getBytes(), untilOffset.toString.getBytes())

    table.put(put)

  }
}
